from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import httpx
import asyncio
import time
import json
import logging
from typing import Dict, Any, Optional, List
from config import config
import os

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("rag-api")

# 创建FastAPI应用
app = FastAPI(
    title="百度智能云RAG问答API",
    description="基于百度智能云千帆平台的知识库问答系统",
    version="1.0.0"
)

# 启用CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源，生产环境应设置为具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 用于缓存access token
token_cache = {
    "access_token": "",
    "expires_at": 0
}

# 修改历史记录的数据结构和处理方式
dialogue_history = []

# 请求模型
class QuestionRequest(BaseModel):
    question: str
    history: Optional[List[Dict[str, str]]] = None
    stream: Optional[bool] = False
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.8
    
# 响应模型
class AnswerResponse(BaseModel):
    answer: str
    references: Optional[List[Dict[str, Any]]] = None
    
# 获取access token的异步函数
async def get_access_token() -> str:
    """
    获取百度智能云access token
    如果缓存的token未过期，则直接返回缓存的token
    否则请求新token并更新缓存
    """
    global token_cache
    
    try:
        # 检查API密钥配置
        if not config.API_KEY or not config.SECRET_KEY:
            logger.error("API密钥未配置")
            raise HTTPException(status_code=500, detail="API密钥未配置，请检查配置文件")
            
        # 检查缓存的token是否有效
        current_time = time.time()
        if token_cache["access_token"] and token_cache["expires_at"] > current_time:
            logger.info("使用缓存的access token")
            return token_cache["access_token"]
        
        logger.info("请求新的access token")
        logger.info(f"使用API_KEY: {config.API_KEY[:8]}...")  # 只打印前8位
        
        # 构建请求URL和参数
        token_url = "https://aip.baidubce.com/oauth/2.0/token"
        params = {
            "grant_type": "client_credentials",
            "client_id": config.API_KEY,
            "client_secret": config.SECRET_KEY
        }
        
        logger.info(f"请求token URL: {token_url}")
        logger.info(f"请求参数: {params}")
        
        # 设置代理（如果需要）
        proxies = None
        if config.HTTP_PROXY:
            proxies = {
                "http://": config.HTTP_PROXY,
                "https://": config.HTTP_PROXY
            }
            logger.info(f"使用代理: {config.HTTP_PROXY}")
        else:
            # 尝试从环境变量获取代理设置
            http_proxy = os.environ.get('HTTP_PROXY') or os.environ.get('http_proxy')
            https_proxy = os.environ.get('HTTPS_PROXY') or os.environ.get('https_proxy')
            if http_proxy or https_proxy:
                proxies = {
                    "http://": http_proxy,
                    "https://": https_proxy or http_proxy
                }
                logger.info(f"使用环境变量代理: HTTP={http_proxy}, HTTPS={https_proxy}")
        
        # 设置详细的客户端选项
        client_options = {
            "timeout": 30.0,  # 增加超时时间
            "verify": config.VERIFY_SSL,  # 使用配置的SSL验证设置
            "proxies": proxies if proxies else None,
            "follow_redirects": True
        }
        
        logger.info(f"客户端选项: {client_options}")
        
        try:
            async with httpx.AsyncClient(**client_options) as client:
                try:
                    response = await client.post(token_url, params=params)
                    logger.info(f"收到响应状态码: {response.status_code}")
                    logger.info(f"响应头: {dict(response.headers)}")
                    
                    # 检查响应状态
                    if response.status_code != 200:
                        error_msg = f"获取access token失败: 状态码={response.status_code}, 响应={response.text}"
                        logger.error(error_msg)
                        raise HTTPException(status_code=response.status_code, detail=error_msg)
                    
                    result = response.json()
                    logger.info("成功获取token响应")
                    logger.info(f"响应内容: {json.dumps(result, ensure_ascii=False)}")
                    
                    if "access_token" not in result:
                        error_msg = f"access token响应格式错误: {result}"
                        logger.error(error_msg)
                        raise HTTPException(status_code=500, detail=error_msg)
                    
                    # 更新token缓存
                    token_cache["access_token"] = result["access_token"]
                    token_cache["expires_at"] = current_time + result.get("expires_in", 2592000) - 600
                    
                    logger.info(f"成功获取新token，有效期至: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(token_cache['expires_at']))}")
                    return token_cache["access_token"]
                    
                except httpx.TimeoutException as e:
                    error_msg = f"请求access token超时: {str(e)}, 错误详情: {e.__dict__}"
                    logger.error(error_msg)
                    raise HTTPException(status_code=504, detail=error_msg)
                except httpx.NetworkError as e:
                    error_msg = f"网络连接错误: {str(e)}, 错误详情: {e.__dict__}"
                    logger.error(error_msg)
                    raise HTTPException(status_code=503, detail=error_msg)
                except httpx.RequestError as e:
                    error_msg = f"请求access token时发生错误: {str(e)}, 错误详情: {e.__dict__}"
                    logger.error(error_msg)
                    raise HTTPException(status_code=503, detail=error_msg)
                except Exception as e:
                    error_msg = f"请求过程中发生未知错误: {str(e)}, 错误详情: {e.__dict__ if hasattr(e, '__dict__') else '无详细信息'}"
                    logger.error(error_msg)
                    raise HTTPException(status_code=500, detail=error_msg)
        except Exception as e:
            error_msg = f"创建HTTP客户端时发生错误: {str(e)}, 错误类型: {type(e)}"
            logger.error(error_msg)
            raise HTTPException(status_code=500, detail=error_msg)
            
    except Exception as e:
        error_msg = f"获取access token时发生未知错误: {str(e)}, 错误类型: {type(e)}"
        logger.error(error_msg)
        raise HTTPException(status_code=500, detail=error_msg)

# 调用RAG接口的异步函数
async def call_rag_api(question: str, history: Optional[List[Dict[str, str]]] = None, 
                       temperature: float = 0.7, top_p: float = 0.8) -> Dict[str, Any]:
    try:
        # 获取access token
        access_token = await get_access_token()
        
        # 设置代理（如果需要）
        proxies = None
        if config.HTTP_PROXY:
            proxies = {
                "http://": config.HTTP_PROXY,
                "https://": config.HTTP_PROXY
            }
            logger.info(f"使用代理: {config.HTTP_PROXY}")
        else:
            # 尝试从环境变量获取代理设置
            http_proxy = os.environ.get('HTTP_PROXY') or os.environ.get('http_proxy')
            https_proxy = os.environ.get('HTTPS_PROXY') or os.environ.get('https_proxy')
            if http_proxy or https_proxy:
                proxies = {
                    "http://": http_proxy,
                    "https://": https_proxy or http_proxy
                }
                logger.info(f"使用环境变量代理: HTTP={http_proxy}, HTTPS={https_proxy}")
        
        # 准备请求数据
        messages = []
        if history:
            for msg in history:
                messages.append(msg)
        
        # 添加当前问题
        messages.append({"role": "user", "content": question})
        
        # 基本payload
        payload = {
            "messages": messages,
            "temperature": temperature,
            "top_p": top_p,
            "model": "ERNIE-Bot-4"
        }
        
        if config.SYSTEM_PROMPT:
            payload["system"] = config.SYSTEM_PROMPT
        
        logger.info(f"发送模型请求: {json.dumps(payload, ensure_ascii=False)}")
        
        # 设置重试次数
        max_retries = 3
        retry_count = 0
        
        while retry_count < max_retries:
            try:
                # 发送请求
                async with httpx.AsyncClient(
                    timeout=60.0,  # 增加超时时间到60秒
                    verify=config.VERIFY_SSL,
                    proxies=proxies if proxies else None,
                    follow_redirects=True
                ) as client:
                    try:
                        response = await client.post(
                            f"{config.RAG_URL}?access_token={access_token}",
                            json=payload,
                            headers={"content-type": "application/json"}
                        )
                        
                        logger.info(f"收到响应状态码: {response.status_code}")
                        logger.info(f"响应头: {dict(response.headers)}")
                        logger.info(f"完整响应: {response.text}")
                        
                        # 检查响应
                        if response.status_code != 200:
                            error_msg = f"模型接口请求失败: 状态码={response.status_code}, 响应={response.text}"
                            logger.error(error_msg)
                            if retry_count < max_retries - 1:
                                retry_count += 1
                                logger.info(f"正在进行第 {retry_count} 次重试...")
                                continue
                            raise HTTPException(status_code=response.status_code, detail=error_msg)
                        
                        result = response.json()
                        logger.info(f"收到模型响应: {json.dumps(result, ensure_ascii=False)}")
                        
                        # 处理响应结果
                        answer = result.get("result", "")
                        if not answer and "response" in result:
                            answer = result.get("response", "")
                        
                        if not answer:
                            error_msg = "模型返回的答案为空"
                            logger.error(error_msg)
                            if retry_count < max_retries - 1:
                                retry_count += 1
                                logger.info(f"正在进行第 {retry_count} 次重试...")
                                continue
                            raise HTTPException(status_code=500, detail=error_msg)
                        
                        return {
                            "answer": answer,
                            "references": []
                        }
                        
                    except httpx.TimeoutException as e:
                        error_msg = f"请求模型接口超时: {str(e)}"
                        logger.error(error_msg)
                        if retry_count < max_retries - 1:
                            retry_count += 1
                            logger.info(f"正在进行第 {retry_count} 次重试...")
                            continue
                        raise HTTPException(status_code=504, detail=error_msg)
                    except httpx.NetworkError as e:
                        error_msg = f"网络连接错误: {str(e)}"
                        logger.error(error_msg)
                        if retry_count < max_retries - 1:
                            retry_count += 1
                            logger.info(f"正在进行第 {retry_count} 次重试...")
                            continue
                        raise HTTPException(status_code=503, detail=error_msg)
                    except Exception as e:
                        error_msg = f"请求过程中发生未知错误: {str(e)}"
                        logger.error(error_msg)
                        if retry_count < max_retries - 1:
                            retry_count += 1
                            logger.info(f"正在进行第 {retry_count} 次重试...")
                            continue
                        raise HTTPException(status_code=500, detail=error_msg)
                        
            except Exception as e:
                error_msg = f"创建HTTP客户端时发生错误: {str(e)}"
                logger.error(error_msg)
                if retry_count < max_retries - 1:
                    retry_count += 1
                    logger.info(f"正在进行第 {retry_count} 次重试...")
                    continue
                raise HTTPException(status_code=500, detail=error_msg)
                
    except Exception as e:
        error_msg = f"调用模型接口时发生错误: {str(e)}"
        logger.error(error_msg)
        if isinstance(e, HTTPException):
            raise e
        raise HTTPException(status_code=500, detail=error_msg)

# API端点 - 健康检查
@app.get("/")
async def root():
    return {"status": "ok", "message": "RAG问答API服务运行正常", "mode": "知识库模式" if config.USE_KNOWLEDGE_BASE else "普通对话模式"}

# API端点 - 测试连接
@app.get("/api/test")
async def test_connection():
    """
    简单的API测试端点，用于测试与后端的连接
    
    Returns:
        包含状态和消息的简单JSON响应
    """
    logger.info("收到连接测试请求")
    
    try:
        # 尝试获取token，测试百度API连接
        token = await get_access_token()
        return {
            "status": "ok", 
            "message": "与Python后端和百度API的连接正常",
            "token_status": "有效" if token else "无效"
        }
    except Exception as e:
        logger.error(f"连接测试失败: {str(e)}")
        return {
            "status": "error",
            "message": f"连接测试失败: {str(e)}",
            "token_status": "无效"
        }

# API端点 - RAG问答
@app.post("/api/ask", response_model=AnswerResponse)
async def ask_question(request: QuestionRequest, background_tasks: BackgroundTasks):
    """
    处理用户问题，返回基于知识库的回答
    
    Args:
        request: 包含问题和可选参数的请求体
        
    Returns:
        包含回答和引用信息的响应
    """
    logger.info(f"收到问题: {request.question}")
    
    try:
        # 调用RAG接口
        result = await call_rag_api(
            question=request.question,
            history=request.history,
            temperature=request.temperature,
            top_p=request.top_p
        )
        
        # 将对话保存到历史记录中
        history_item = {
            "id": f"hist_{int(time.time())}",
            "question": request.question,
            "answer": result["answer"],
            "references": result["references"],
            "time": int(time.time())
        }
        
        # 添加到历史记录
        dialogue_history.insert(0, history_item)
        
        # 限制历史记录数量为最新的50条
        if len(dialogue_history) > 50:
            dialogue_history.pop()
        
        # 返回结果
        return AnswerResponse(
            answer=result["answer"],
            references=result["references"]
        )
        
    except Exception as e:
        logger.error(f"处理问题时发生错误: {str(e)}")
        # 如果是HTTPException则继续抛出，否则包装为500错误
        if isinstance(e, HTTPException):
            raise e
        raise HTTPException(status_code=500, detail=f"内部服务器错误: {str(e)}")

# API端点 - 获取历史记录
@app.get("/api/history")
async def get_history(sort_type: str = "new"):
    """
    获取对话历史记录
    
    Args:
        sort_type: 排序方式，new-最新在前，old-最早在前
        
    Returns:
        包含历史记录列表的响应
    """
    try:
        logger.info(f"获取历史记录，排序方式: {sort_type}")
        
        # 复制历史记录列表，避免直接修改原始数据
        sorted_history = dialogue_history.copy()
        
        # 根据排序类型排序
        if sort_type == "old":
            sorted_history.sort(key=lambda x: x.get("time", 0))
        else:  # 默认按最新排序
            sorted_history.sort(key=lambda x: x.get("time", 0), reverse=True)
            
        return {
            "code": 200,
            "msg": "获取历史记录成功",
            "data": {
                "list": sorted_history,
                "total": len(sorted_history)
            }
        }
    except Exception as e:
        error_msg = f"获取历史记录失败: {str(e)}"
        logger.error(error_msg)
        return {
            "code": 500,
            "msg": "获取历史记录失败",
            "data": {
                "list": [],
                "total": 0
            }
        }

# API端点 - 添加历史记录
@app.post("/api/history/add")
async def add_history(history_item: dict):
    """
    添加一条对话历史记录
    
    Args:
        history_item: 包含问题和答案的字典
    """
    try:
        logger.info(f"添加历史记录: {json.dumps(history_item, ensure_ascii=False)}")
        
        # 确保必要的字段存在
        if "question" not in history_item or "answer" not in history_item:
            raise ValueError("历史记录必须包含问题和答案")
            
        # 添加时间戳和ID（如果没有）
        if "time" not in history_item:
            history_item["time"] = int(time.time() * 1000)  # 使用毫秒时间戳
        if "id" not in history_item:
            history_item["id"] = f"hist_{history_item['time']}"
            
        # 检查是否已存在相同的问题（最近5分钟内）
        current_time = history_item["time"]
        five_minutes_ago = current_time - (5 * 60 * 1000)  # 5分钟前的时间戳
        
        # 查找最近5分钟内的相同问题
        recent_same_question = None
        for item in dialogue_history:
            if (item["question"] == history_item["question"] and 
                item["time"] > five_minutes_ago):
                recent_same_question = item
                break
        
        # 如果找到相同的问题，更新它而不是添加新记录
        if recent_same_question:
            recent_same_question.update(history_item)
            logger.info("更新已存在的历史记录")
        else:
            # 添加到历史记录
            dialogue_history.insert(0, history_item)
            logger.info("添加新的历史记录")
            
        # 限制历史记录数量为最新的50条
        if len(dialogue_history) > 50:
            dialogue_history.pop()
            
        return {
            "code": 200,
            "msg": "添加历史记录成功",
            "data": history_item
        }
    except Exception as e:
        error_msg = f"添加历史记录失败: {str(e)}"
        logger.error(error_msg)
        return {
            "code": 500,
            "msg": f"添加历史记录失败: {str(e)}",
            "data": None
        }

# API端点 - 清空历史记录
@app.delete("/api/history")
async def clear_history():
    """
    清空所有对话历史记录
    """
    try:
        logger.info("清空历史记录")
        dialogue_history.clear()
        return {
            "code": 200,
            "msg": "清空历史记录成功",
            "data": None
        }
    except Exception as e:
        error_msg = f"清空历史记录失败: {str(e)}"
        logger.error(error_msg)
        return {
            "code": 500,
            "msg": f"清空历史记录失败: {str(e)}",
            "data": None
        }

# 启动服务的命令行入口
if __name__ == "__main__":
    import uvicorn
    logger.info(f"启动服务，模式: {'知识库模式' if config.USE_KNOWLEDGE_BASE else '普通对话模式'}")
    logger.info(f"使用API地址: {config.RAG_URL}")
    
    uvicorn.run(
        "main:app", 
        host=config.HOST, 
        port=config.PORT, 
        reload=config.DEBUG
    ) 